#pragma once
#include "ast.h"
#include "symbol_table.h"
#include <variant>
#include <vector>
#include <string>
#include <unordered_map>

namespace ir
{
    /* 文法 */
    // FunDefs -> FunDef FunDefs
    // FunDef -> "fun" SYMBOL "(" FunParams? ")" (": i32")? "{" FunBody "}"
    // FunParams -> FunParam ("," FunParams)?
    // FunParam -> SYMBOL ": i32"
    // FunBody -> Blocks?
    // Blocks -> Block Blocks?
    // Block -> SYMBOL ":" Statements? EndStatement
    // Statements -> Statement Statements?
    // Statement -> SymbolDef | Store | FunCall
    // SymbolDef -> SYMBOL "=" (MemoryDeclaration | Load | Mv | BinaryExpr | FunCall)
    // MemoryDeclaration -> "alloc i32"
    // Load -> "load" SYMBOL
    // Mv -> "mv" Value
    // BinaryExpr -> BINARY_OP Value "," Value
    // Store -> "store" Value "," SYMBOL
    // FunCall -> "call" SYMBOL "(" FunArgs? ")"
    // FunArgs -> Value ("," FunArgs)?
    // EndStatement -> Branch | Jump | Return
    // Branch -> "br" Value "," SYMBOL "," SYMBOL
    // Jump -> "jmp" SYMBOL
    // Return -> "ret" Value?
    // Value -> SYMBOL | INT
    
    /* SSA */
    // 具有SSA形式：所有变量只允许被赋值一次
    // 但不严格：允许通过store指令来修改栈上的变量

    /* 存储 */
    // IR -> Function -> (Basic) Block -> Instruction

    /* TODO */
    // 假定测试用例语义正确，跳过的部分语义分析的工作
    // 1. 调用函数时检查函数的返回值类型和参数个数
    // 2. 验证返回值为int的函数是否所有路径都有return语句

    typedef std::variant<std::string, int> Operand;

    enum Type { TYPE_INT, TYPE_VOID };

    enum Opcode
    {
        OP_ADD, OP_SUB, OP_MUL, OP_DIV, OP_MOD,
        OP_NE, OP_EQ, OP_GT, OP_LT, OP_GE, OP_LE,
        OP_AND, OP_OR, OP_XOR, OP_SHL, OP_SHR, OP_SAR,
        OP_LOAD, OP_STORE, OP_ALLOC, OP_MV,
        OP_BR, OP_JMP, OP_RET, OP_CALL
    };
    
    struct Instruction
    {
        std::string ToString();
        Opcode op;
        std::string dest;
        std::vector<Operand> srcs;
    };

    struct Block
    {
        std::string ToString();
        std::string label;
        std::vector<Instruction> instrs;
    };

    struct Function
    {
        std::string ToString();
        std::string name;
        std::vector<std::string> params;
        Type type;
        std::vector<Instruction> alloc_instrs;
        std::vector<Block> blocks;
    };

    struct IR
    {
        std::string ToString();
        std::vector<Function> funcs;
    };

    class IRGenerator 
    {
    public:
        IR Generate(ast::CompUnitAST& ast); 
    private:
        // 访问不同ast结点的操作
        void Visit(ast::BaseAST& node);
        void Visit(ast::ListAST& node);
        void Visit(ast::CompUnitAST& node);
        void Visit(ast::FuncDefAST& node);
        void Visit(ast::ParamAST& node);
        void Visit(ast::StmtAST& node);
        void Visit(ast::AssignStmtAST& node);
        void Visit(ast::DefStmtAST& node);
        void Visit(ast::IfStmtAST& node);
        void Visit(ast::WhileStmtAST& node);
        void Visit(ast::ExprAST& node);
        void Visit(ast::BinaryExprAST& node);
        void Visit(ast::UnaryExprAST& node);
        void Visit(ast::CallExprAST& node);
        void Visit(ast::ArgAST& node);
        void Visit(ast::IdAST& node);
        void Visit(ast::NumAST& node);

        // 其他辅助函数
        void VisitBlockStmt(ast::StmtAST& node);
        void VisitBreakStmt();
        void VisitContStmt();
        void VisitRetStmt(ast::StmtAST& node);
        void VisitBinaryExpr(ast::BinaryExprAST& node);
        void VisitShortCircuitExpr(ast::BinaryExprAST& node);
        void PushFunc();
        void PushParam();
        void PushBlock();
        void PushInstr();
        void PushOperand();
        void PushCallOperand();
        void PushBinaryInstr();
        void PushLoadInstr();
        void PushStoreInstr();
        void PushAllocInstr();
        void PushBrInstr();
        void PushJmpInstr();
        void PushRetInstr();
        void PushCallInstr();
        std::string GetTempSymbol();
        std::string GetNamedSymbol(std::string& name);
        
        // 从ast读取的数据
        IR ir;
        Function func;
        Block block;
        Instruction instr;
        Opcode op;
        std::string id;
        Operand operand1, operand2, operand3;
        std::vector<Operand> call_operands;

        // 其他辅助变量
        symbol_table::SymbolTable<int> symbol_table; 
        bool is_rval;
        bool is_in_block_stmt;
        int local_symbol_cnt;
        std::vector<std::string> while_entry_labels;
        std::vector<std::string> while_end_labels;
    };
    
}